#include "tensorflowCCVisionLib.h"
// #define DEBUG_INFO 
#include<ctime>
TensorflowCCVisoin::TensorflowCCVisoin(void){
    input_tensor_name="image_tensor";
    out_put_nodes.reserve(3);
    out_put_nodes.push_back("detection_scores");  //detection_scores  detection_classes  detection_boxes
    out_put_nodes.push_back("detection_classes");
    out_put_nodes.push_back("detection_boxes");
}

TensorflowCCVisoin::~TensorflowCCVisoin(){
    if(session != nullptr )
        session->Close();
    graphdef.Clear();

}
/**
 * 
 */
int TensorflowCCVisoin::InitParam( int img_width, int img_height){


    input_width = img_width;

    input_height = img_height;
    // }

    Status status = NewSession(SessionOptions(), &session);
    if (!status.ok()) {
        cout << "ERROR: NewSession failed..." << std::endl;
        cout << status.ToString() << "\n";
        return -1;
    }
    
    //初始化 tensor 向量
    resized_tensor = Tensor(DT_UINT8, TensorShape({1,input_height,input_width,3})); //DT_FLOAT

    return 0;
}



/**
 * 
 */
int TensorflowCCVisoin::LoadImg(cv::Mat tempImage){
    // resized_tensor.~Tensor();

    return  CVMat_to_Tensor(tempImage, &resized_tensor, input_height, input_width);
    
}
/**
 * 
 */
int TensorflowCCVisoin::LoadPbParam(){
    Status status_load = ReadBinaryProto(Env::Default(), model_path, &graphdef); //从pb文件中读取图模型;
    if (!status_load.ok()) {
        cout << "ERROR: Loading model failed..." << model_path << std::endl;
        cout << status_load.ToString() << "\n";
        return -1;
    }
    //这里需要待修改
    Status status_create = session->Create(graphdef); //将模型导入会话Session中;
    if (!status_create.ok()) {
        cout << "ERROR: Creating graph in session failed..." << status_create.ToString() << std::endl;
        return -1;
    }
    cout << "<----Successfully created session and load graph.------->"<< endl;
    return 0;
}
 
/**
 * 
 */
int TensorflowCCVisoin::TfDetection(){

    // outputs.clear();
    clock_t start, endti;
    start = clock();     
    Status status_run = session->Run({{input_tensor_name, resized_tensor}}, {out_put_nodes}, {}, &outputs);
    endti = clock(); 
    double endtime=(double)(endti-start)/CLOCKS_PER_SEC;
    cout << "time "<< endtime<<" s "<<endl;

    if (!status_run.ok()) {
        cout << "ERROR: RUN failed..."  << std::endl;
        cout << status_run.ToString() << "\n";
        return -1;
    }
#ifdef DEBUG_INFO
    //把输出值给提取出
    cout << "Output tensor size:" << outputs.size() << std::endl;  //3
    for (int i = 0; i < outputs.size(); i++)
    {
        cout << outputs[i].DebugString()<<" "<< std::endl;   // [1, 50], [1, 50], [1, 50, 4]
    }
#endif

    return 0;
    // cvtColor(img, img, CV_RGB2BGR);  // opencv读入的是BGR格式输入网络前转为RGB
    // resize(img,img,cv::Size(480,280));  // 模型输入图像大小
   
}

int TensorflowCCVisoin::TfGetResult(std::vector<float> &result){
     
    int pre_num = outputs[0].dim_size(1);  // 50  模型预测的目标数量
    auto tmap_pro = outputs[0].tensor<float, 2>();  //第一个是score输出shape为[1,50]
    auto tmap_clas = outputs[1].tensor<float, 2>();  //第二个是class输出shape为[1,50]
    auto tmap_coor = outputs[2].tensor<float, 3>();  //第三个是coordinate输出shape为[1,50,4]
    float probability = 0.4;  //自己设定的score阈值
    pre_num = 1 ;

    for (int pre_i = 0; pre_i < pre_num; pre_i++)
    {
        if (tmap_pro(0, pre_i) < probability)
        {
            cout<<" score is "<< tmap_pro(0, pre_i)<<endl;
            return -1;
        }

        string id = to_string(int(tmap_clas(0, pre_i)));
        int xmin = int(tmap_coor(0, pre_i, 1) * input_width);
        int ymin = int(tmap_coor(0, pre_i, 0) * input_height);
        int xmax = int(tmap_coor(0, pre_i, 3) * input_width);
        int ymax = int(tmap_coor(0, pre_i, 2) * input_height);

        cout << "Class ID: " << tmap_clas(0, pre_i) << endl;
        cout << "Probability: " << tmap_pro(0, pre_i) << endl;
        cout << "Xmin is: " << xmin << endl;
        cout << "Ymin is: " << ymin << endl;
        cout << "Xmax is: " << xmax << endl;
        cout << "Ymax is: " << ymax << endl;

        result.reserve(6);
        result.push_back(tmap_clas(0, pre_i)); result.push_back(tmap_pro(0, pre_i));
        result.push_back(xmin); result.push_back(ymin);
        result.push_back(xmax);result.push_back(ymax);

        break;
        // rectangle(img, cvPoint(xmin, ymin), cvPoint(xmax, ymax), Scalar(255, 0, 0), 1, 1, 0);
        // putText(img, id, cvPoint(xmin, ymin), FONT_HERSHEY_COMPLEX, 1.0, Scalar(255,0,0), 1);
    }

    return 0;
}


 
// // 定义一个函数讲OpenCV的Mat数据转化为tensor，python里面只要对cv2.read读进来的矩阵进行np.reshape之后，
// // 数据类型就成了一个tensor，即tensor与矩阵一样，然后就可以输入到网络的入口了，但是C++版本，我们网络开放的入口
// // 也需要将输入图片转化成一个tensor，所以如果用OpenCV读取图片的话，就是一个Mat，然后就要考虑怎么将Mat转化为
// // Tensor了
int TensorflowCCVisoin::CVMat_to_Tensor(Mat img,Tensor* output_tensor,int input_rows,int input_cols)
{

    assert(img.cols == input_cols && img.rows == input_rows);
    //归一化
    img.convertTo(img,CV_8UC3);  // CV_32FC3
    // img=1-img/255;
 
    //创建一个指向tensor的内容的指针
    uint8 *p = output_tensor->flat<uint8>().data();
 
    //创建一个Mat，与tensor的指针绑定,改变这个Mat的值，就相当于改变tensor的值
    tempMat = cv::Mat(input_rows, input_cols, CV_8UC3, p);
    img.convertTo(tempMat,CV_8UC3);  //keypoint
  
    return 0;
}

int TensorflowCCVisoin::init( int img_width, int img_height){
    int rtn = InitParam( img_width, img_height);
    if( rtn != 0 ){
        return -1;
    }

    return 0;

}

int TensorflowCCVisoin::LoadModel(const string& model_path){
    if(model_path.size() < 2 ){
        this->model_path = "/home/de/pbmodel/0704-10000/frozen_inference_graph.pb";
    }else{
        this->model_path = model_path;

    }

    int rtn = LoadPbParam();
    if( rtn != 0){
        return -2;
    }

    return 0;
}

/**
 * 识别跟踪 这里会自动resize 图像的大小
 */
int TensorflowCCVisoin::detection(cv::Mat tempImage, vector<float> &pos){
    int rtn = LoadImg(tempImage);
    if( rtn !=0 ){
        return -1;
    }

    rtn = TfDetection();
    if( rtn !=0 ){
        return -2;
    }
    rtn = TfGetResult(pos);
    if( rtn !=0 ){
        return -3;
    }
    return 0;
} 